"""从一个大的 reflow dataset 中构建一个小的 subset , 随机选取 samples
"""
import sys
sys.path.append('.')
from reflow.data.dataset import get_reflow_dataset
from reflow.data.utils import LMDB_ndarray, data2lmdb
from reflow.utils import set_seed
import random
from tqdm.auto import tqdm
from pathlib import Path
import numpy as np
from loguru import logger
import json
from argparse import Namespace

if __name__ == "__main__":
    args=Namespace(
        src_ds_path='data/coco2014_reflow/alt_gen_val10k' ,
        tgt_ds_path='data/coco2014_reflow/alt_gen_val_rnd20',
        num_samples=20,
        seed=32,
        type='random' # random, head
    )
    set_seed(seed=args.seed)
    
    logger.info(f'use random seed {args.seed}')
    logger.info(f'data saved in <<{args.tgt_ds_path}>>')
    
    save_dir = Path(args.tgt_ds_path)
    save_dir.mkdir(parents=True, exist_ok=True)
    content_dir = save_dir / 'content'
    content_dir.mkdir(parents=True, exist_ok=True)
    image_dir = content_dir / 'images'
    image_dir.mkdir(parents=True, exist_ok=True)
    
    # 准备保存每个 (noise,latent) pair , 以及所有的 captions
    
    src_ds=get_reflow_dataset(
        data_root=args.src_ds_path, 
        src_type='lmdb',
    )
    if args.type=='random':
        rnd_indices=[random.randint(0,len(src_ds)-1) for _ in range(args.num_samples)]
    elif args.type=='head':
        rnd_indices=list(range(args.num_samples))

    caps=[]
    for cnt, i in enumerate(tqdm(rnd_indices)):
        data=src_ds[i]
        caps.append(data['caption'])
        pair = np.stack((data['noise'], data['latent']), axis=0)
        np.save(str(image_dir / f'{cnt}.npy'), pair)

    with open(str(content_dir / 'captions.txt'), 'w') as f:
        f.write('\n'.join(caps))
        
    data2lmdb(str(image_dir))
    
    json.dump(vars(args), open(str(save_dir / 'index.json'), 'w'))